import json
from collections import deque
from text2image import text2image
from text2audio import text2audio
from text2text import text2text
from text2text import text2text_deepseek_7b_math
from text2text import text2text_qwen_7b_coder
from text2text import text2text_internlm2_5_7b_chat
from text2video import text2video


def text2image_prompt_generation(info,target):
    prompt = f"""
    <info>{info}</info>
    <target>{target}</target>
    你是AI绘画专家，善于使用midjourney和stable diffusion等工具绘画。
    你的目标是画出target,必要的参考信息是info。
    现在你要构造一个midjourney风格的绘画提示词。
    直接输出绘画提示词。
    要求，80字左右。中文。
    """
    return text2text(prompt)
    pass

def text2audio_prompt_generation(info,target):
    prompt = f"""
    <info>{info}</info>
    <target>{target}</target>
    你是人类，善于用各种语言风格说话。
    你的目标是target,必要的参考信息是info，你以人类的口吻，自主设计一个语言风格，输出一个演讲文案。
    """
    return text2text(prompt)
    
    pass


def text2video_prompt_generation(info,target):
    prompt = f"""
    <info>{info}</info>
    <target>{target}</target>
    你是视频设计师，现在要设计一个可以生成视频的提示词。
    提示词公式 = 主体（主体描述）+运动+场景（场景描述）+（镜头语言+光影+氛围），括号里的内容可选填。
    以上公式最核心的构成就是主体+运动+场景，这也是描述一个视频画面最简单、最基本的单元。
    你的目标是target，必要的参考信息是info。
    然后根据提示词公式，直接输出一个视频描述。100字。中文。
    """
    return text2text(prompt)
    pass

class Task:
    def __init__(self, task_id, task_name, task_type, content, dependencies):
        self.id = task_id
        self.task_name = task_name
        self.type = task_type
        self.content = content
        self.dependencies = dependencies
        self.result = None

def topological_sort(tasks):
    adjacency = {task.id: [] for task in tasks}
    in_degree = {task.id: 0 for task in tasks}

    for task in tasks:
        for dep_id in task.dependencies:
            adjacency[dep_id].append(task.id)
            in_degree[task.id] += 1

    queue = deque([task.id for task in tasks if in_degree[task.id] == 0])
    sorted_order = []

    while queue:
        task_id = queue.popleft()
        sorted_order.append(task_id)
        for next_task_id in adjacency[task_id]:
            in_degree[next_task_id] -= 1
            if in_degree[next_task_id] == 0:
                queue.append(next_task_id)

    if len(sorted_order) != len(tasks):
        raise ValueError("存在循环依赖，无法进行拓扑排序")

    return sorted_order

def text_generation_task(task, input_data):
    info = ";".join(input_data)
    target = task.content

    prompt = f"""
    {info},
    {target}
    """
    
    model_select_prompt = f"""

    不同的函数善于处理不同的问题。
    ## 候选函数名:
    text2text,
    text2text_qwen_7b_coder,
    text2text_deepseek_7b_math,
    text2text_internlm2_5_7b_chat
    ## 擅于处理的问题类型
    text2text:擅于创作，内容生成，强大的知识库
    text2text_qwen_7b_coder:擅于编写代码，编程
    text2text_deepseek_7b_math:擅于处理数学问题
    text2text_internlm2_5_7b_chat:擅长英语，擅长当个客服，擅长处理商品问题。智能客服，擅长处理顾客问题。
    ##
    <text>{prompt}</text>
    现在分析和识别<text>中的问题类别，从候选函数名中给出一个最适合处理它的函数。
    直接且输出函数名。
    
    输出例子:
    text2text_qwen_7b_coder

    """

    fun_name = text2text(model_select_prompt)
    text2text_fun_ditc = {
        "text2text":text2text,
        "text2text_qwen_7b_coder":text2text_qwen_7b_coder,
        "text2text_deepseek_7b_math":text2text_deepseek_7b_math,
        "text2text_internlm2_5_7b_chat":text2text_internlm2_5_7b_chat
    }

    text2text_fun = text2text_fun_ditc[fun_name]
    
    
    return text2text_fun(prompt)

def image_generation_task(task, input_data):
    info = ";".join(input_data)
    target = task.content

    image_prompt = text2image_prompt_generation(info,target)

    return text2image(image_prompt)

def audio_generation_task(task, input_data):
    info = ";".join(input_data)
    target = task.content
    
    audio_prompt = text2audio_prompt_generation(info,target)
    return text2audio(audio_prompt)

def video_generation_task(task, input_data):
    info = ";".join(input_data)
    target = task.content
    
    audio_prompt = text2video_prompt_generation(info,target)
    return text2video(audio_prompt)

def execute_task(task, input_data):
    """根据任务类型调用相应的函数"""
    task_type_map = {
        "文生文": text_generation_task,
        "文生图": image_generation_task,
        "文生音频": audio_generation_task,
        "文生视频": video_generation_task
    }
    
    func = task_type_map.get(task.type)
    if func:
        return func(task, input_data)
    else:
        return f"未知任务类型：{task.type}"

def main(json_data):
    data = json_data
    tasks = data["tasks"]

    task_dict = {}
    task_list = []
    for task in tasks:
        task_obj = Task(
            task["id"],
            task["task_name"],
            task["task_type"],
            task["content"],
            task["dependencies"]
        )
        task_dict[task_obj.id] = task_obj
        task_list.append(task_obj)

    sorted_order = topological_sort(task_list)
    print("拓扑排序后的任务执行顺序：", sorted_order)

    for task_id in sorted_order:
        current_task = task_dict[task_id]
        print(f"\n执行任务：{current_task.id} - {current_task.task_name}")
        
        input_data = [
            task_dict[dep_id].result
            for dep_id in current_task.dependencies
        ]
        
        current_task.result = execute_task(current_task, input_data)
        print(f"任务执行结果：{current_task.result}")

    print("\n所有任务执行完成，结果如下：")
    results = []
    for task_id in sorted_order:
        results.append({
            "id": task_id,
            "task_name": task_dict[task_id].task_name,
            "task_type": task_dict[task_id].type,  # 添加任务类型
            "content": task_dict[task_id].content,  # 添加任务内容
            "dependencies": task_dict[task_id].dependencies,  # 添加任务依赖
            "result": task_dict[task_id].result
        })
        print(f"任务 {task_id} 的结果：{task_dict[task_id].result}")

    # 将结果保存为 JSON 文件
    with open("task_results.json", "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=4)
    print("\n结果已保存到 task_results.json 文件中。")

    return results
if __name__ == "__main__":
    # 打开 JSON 文件并读取数据
    with open("tasks.json", "r", encoding="utf-8") as file:
        json_data = json.load(file)
    main(json_data)